import face_recognition
import psycopg2
import numpy as np

# 连接 PostgreSQL 数据库
def connect_to_db():
    return psycopg2.connect(
        dbname="postgres",
        user="postgres",
        password="postgres",
        host="192.168.31.171",
        port="5432"
    )

# 提取人脸特征并存储到数据库
def store_face(image_path, name):
    # 加载图片并提取特征
    image = face_recognition.load_image_file(image_path)
    face_encodings = face_recognition.face_encodings(image)

    if len(face_encodings) == 0:
        print(f"未检测到人脸: {image_path}")
        return

    face_encoding = face_encodings[0]  # 取第一个检测到的人脸
    face_encoding_list = face_encoding.tolist()  # 转换为 Python 列表

    # 存储到数据库
    conn = connect_to_db()
    cursor = conn.cursor()
    cursor.execute(
        "INSERT INTO faces (name, encoding) VALUES (%s, %s)",
        (name, face_encoding_list)
    )
    conn.commit()
    cursor.close()
    conn.close()
    print(f"人脸数据已存储: {name}")

# 示例：存储多张人脸
# store_face("../image/jack_ma1.jpeg", "Jack Ma")
# store_face("../image/liuqiangdong1.jpeg", "Liu Qiangdong")

# 检索匹配的人脸
def search_face(image_path):
    # 加载输入图片并提取特征
    input_image = face_recognition.load_image_file(image_path)
    input_face_encodings = face_recognition.face_encodings(input_image)

    if len(input_face_encodings) == 0:
        print("未检测到人脸")
        return

    input_face_encoding = input_face_encodings[0]
    input_face_encoding_list = input_face_encoding.tolist()  # 转换为 Python 列表

    # 查询数据库，计算距离
    conn = connect_to_db()
    cursor = conn.cursor()

    # 使用 L2 距离（欧氏距离）进行排序
    query = """
    SELECT id, name, encoding <-> %s::vector AS distance
    FROM faces
    ORDER BY distance ASC
    LIMIT 1;
    """
    cursor.execute(query, (input_face_encoding_list,))
    result = cursor.fetchone()

    cursor.close()
    conn.close()

    # 输出结果
    threshold = 0.6  # 相似度阈值
    if result and result[2] < threshold:
        print(f"匹配到: {result[1]} (ID: {result[0]}), 相似度分数: {result[2]:.2f}")
    else:
        print("未找到匹配")

# 示例：检索人脸
search_face("../image/liuqiangdong2.jpeg")